"""
数据库配置模块

提供PostgreSQL和TimescaleDB的连接配置和管理功能。
"""

from typing import AsyncGenerator, Optional
from sqlalchemy.ext.asyncio import AsyncSession, create_async_engine, async_sessionmaker
from sqlalchemy.orm import DeclarativeBase
from sqlalchemy import MetaData
from .settings import settings


class Base(DeclarativeBase):
    """SQLAlchemy基础模型类"""
    metadata = MetaData(
        naming_convention={
            "ix": "ix_%(column_0_label)s",
            "uq": "uq_%(table_name)s_%(column_0_name)s",
            "ck": "ck_%(table_name)s_%(constraint_name)s",
            "fk": "fk_%(table_name)s_%(column_0_name)s_%(referred_table_name)s",
            "pk": "pk_%(table_name)s"
        }
    )


class DatabaseConfig:
    """数据库配置类"""
    
    def __init__(self):
        """初始化数据库配置"""
        # 主数据库引擎
        self.engine = create_async_engine(
            settings.database_url,
            pool_size=settings.database_pool_size,
            max_overflow=settings.database_max_overflow,
            pool_pre_ping=True,
            echo=settings.debug,
        )
        
        # 时序数据库引擎
        self.timescale_engine = create_async_engine(
            settings.timescale_url,
            pool_size=settings.database_pool_size,
            max_overflow=settings.database_max_overflow,
            pool_pre_ping=True,
            echo=settings.debug,
        )
        
        # 会话工厂
        self.async_session_maker = async_sessionmaker(
            bind=self.engine,
            class_=AsyncSession,
            expire_on_commit=False,
        )
        
        # 时序数据库会话工厂
        self.timescale_session_maker = async_sessionmaker(
            bind=self.timescale_engine,
            class_=AsyncSession,
            expire_on_commit=False,
        )
    
    async def get_session(self) -> AsyncGenerator[AsyncSession, None]:
        """获取主数据库会话"""
        async with self.async_session_maker() as session:
            try:
                yield session
                await session.commit()
            except Exception:
                await session.rollback()
                raise
            finally:
                await session.close()
    
    async def get_timescale_session(self) -> AsyncGenerator[AsyncSession, None]:
        """获取时序数据库会话"""
        async with self.timescale_session_maker() as session:
            try:
                yield session
                await session.commit()
            except Exception:
                await session.rollback()
                raise
            finally:
                await session.close()
    
    async def create_tables(self):
        """创建数据库表"""
        async with self.engine.begin() as conn:
            await conn.run_sync(Base.metadata.create_all)
    
    async def drop_tables(self):
        """删除数据库表"""
        async with self.engine.begin() as conn:
            await conn.run_sync(Base.metadata.drop_all)
    
    async def close(self):
        """关闭数据库连接"""
        await self.engine.dispose()
        await self.timescale_engine.dispose()


# 全局数据库配置实例
db_config = DatabaseConfig()


async def init_database():
    """初始化数据库"""
    await db_config.create_tables()


async def close_database():
    """关闭数据库连接"""
    await db_config.close()


# 依赖注入函数
async def get_db() -> AsyncGenerator[AsyncSession, None]:
    """FastAPI依赖注入：获取数据库会话"""
    async for session in db_config.get_session():
        yield session


async def get_timescale_db() -> AsyncGenerator[AsyncSession, None]:
    """FastAPI依赖注入：获取时序数据库会话"""
    async for session in db_config.get_timescale_session():
        yield session